# task13_env.py
# 非配平增稳起飞专用环境（分地面/空中两阶段）
import math
import logging
import numpy as np
from gym import spaces
from getDCSdata import get_dcs_data
from task1_env0329 import TakeoffEnv, TARGET_VA
from vjoy_control_new import activate_dcs_window


class StabilizedTakeoffEnv(TakeoffEnv):
    """
    改进后的非配平增稳起飞环境，添加了侧风扰动和湍流强度模拟等因素，
    使其更接近真实飞行环境下的起飞情况。
    """
    def __init__(self):
        try:
            activate_dcs_window()
        except Exception as e:
            logging.warning(f"子类窗口激活失败: {e}")

        super().__init__()

        # 保存父类原始 observation_space
        self._base_obs_space = self.observation_space
        base_dim = self._base_obs_space.shape[0]

        # 角速度归一化(3)(删除) + 侧风信息(3) + 湍流信息(3)
        extra_dim = 6
        # 一次性定义扩展后的状态空间
        self.observation_space = spaces.Box(
            low=-np.inf,
            high=np.inf,
            shape=(base_dim + extra_dim,),
            dtype=np.float32
        )

        # 新增：侧风与湍流参数（用于模拟真实起飞环境中的外界气流扰动）
        self.crosswind_speed = 0.0  # 侧风速度（单位：m/s）
        self.crosswind_direction = 0.0  # 侧风方向（角度，0-360度）
        self.turbulence_intensity = 0.0  # 湍流强度（用于控制噪声标准差，单位与后续添加噪声相关）

        # 动态归一化角速度参数（针对角速度部分）
        # self.omega_window = np.zeros((100, 3))
        # self.current_omega_max = np.array([0.5, 0.5, 0.5])

        # 安全参数
        self.stall_speed = 45.0  # 失速速度（m/s）
        self.max_pitch_ground = 15.0  # 地面阶段最大允许俯仰角（度）

        # 奖励参数
        self.reward_weights = {
            'stability': 0.3,
            'speed': 0.2,
            'altitude': 0.05,
            'distance': 0.05,
            'safety': 0.3,
            'aoa': 0.1
        }

    def reset(self):
        # 1) 暂时切回原维数，避免基类 reset 时断言失败
        saved_space = self.observation_space
        self.observation_space = self._base_obs_space

        # 2) 调用父类 reset()，它会返回原维数的 state
        base_state = super().reset()

        # 3) 恢复成现维数
        self.observation_space = saved_space

        """
        自适应扰动：根据训练进度对角速度施加随机扰动，对滚转和俯仰方向扰动标准差0.5+0.5*progress，初期0.5rad/s，末期达1rad/s
        有助于模拟非配平状态下的初始不稳定性。progress越大，扰动越大
        """
        # 自适应扰动：对角速度施加扰动（保持原有功能）
        """
        progress = min(self.steps / 1e5, 1.0)  # 假设总训练步数为1e5
        rad2deg = 180.0 / math.pi
        perturbation = [
            np.random.normal(0, 0.5 + 0.5 * progress) * rad2deg,  # 滚转扰动（rad/s）
            np.random.normal(0, 0.5 + 0.5 * progress) * rad2deg,  # 俯仰扰动（rad/s）
            np.random.uniform(-0.2 * rad2deg, 0.2 * rad2deg)  # 偏航扰动（rad/s）
        ]
        base_state[24:27] = perturbation
        """

        # 新增：初始化侧风和湍流参数，每个回合开始时随机设置
        self.crosswind_speed = np.random.uniform(0, 10)  # 侧风速度范围：0 ~ 10 m/s
        self.crosswind_direction = np.random.uniform(0, 360)  # 侧风方向：随机0~360度
        self.turbulence_intensity = np.random.uniform(0, 2)  # 湍流强度范围：0 ~ 2（作为标准差）

        # 新增为俯仰（pitch，state[3]）和偏航（yaw，state[2]）添加预设偏差 ±10° 若内部角度单位为度，可直接操作；若为弧度，则需转换（例如：np.deg2rad(10)）
        # preset_offset_deg = np.deg2rad(10)
        preset_offset_deg = 10.0
        # 随机选择偏正或偏负
        # pitch_offset = np.deg2rad(preset_offset_deg) * choice([-1,1]), yaw_offset = np.deg2rad(preset_offset_deg) * choice([-1,1])
        pitch_offset = preset_offset_deg * np.random.choice([-1, 1])
        yaw_offset = preset_offset_deg * np.random.choice([-1, 1])
        # 若内部存储单位为弧度，则应改为 np.deg2rad(preset_offset_deg) * choice([-1,1])
        base_state[3] += pitch_offset
        base_state[2] += yaw_offset

        # 生成扩展后的状态
        augment_state = self._augment_state(base_state)
        self.state = augment_state
        return augment_state

    def step(self, action):
        # 调用父类step
        # super().step()
        saved_space = self.observation_space
        self.observation_space = self._base_obs_space
        """
                if self.is_on_ground:
            # 地面阶段固定襟翼控制,油门、俯仰、滚转、襟翼、扰流板、脚舵
            full_action = np.concatenate([action, [0.0, 0.0]])
        else:
            # 空中阶段重新排列动作次序（[throttle, pitch, airbrake, roll]，并补全扰流板动作，襟翼、脚舵）
            full_action = np.concatenate([action[:2], action[3:4], action[2:3], [0.0, 0.0]])

        next_state, reward, done, info = super().step(full_action)
        """
        # 直接使用6维动作，不做扩充
        next_state, reward, done, info = super().step(action)
        # 保存原始下一个状态
        # next_state = next_base
        # 3) 恢复扩展后的 obs space（46 维）
        self.observation_space = saved_space

        # 安全性检查：如果检测到危险状态，则惩罚并结束回合
        if self._check_dangerous_state(next_state):
            reward -= 1000.0
            done = True

        # 增稳环境奖励：加入额外奖励修正（考虑侧风与湍流）
        reward += self._calculate_balanced_reward(next_state)

        # 将环境扰动（侧风、湍流）作用于状态，使观测更真实
        disturbed = self._apply_environment_disturbance(next_state.copy())
        # next_state = self._apply_environment_disturbance(next_state)

        # 生成扩展状态并保存
        augment_state = self._augment_state(disturbed)
        self.state = augment_state
        return augment_state,reward, done, info
        # return self._augment_state(next_state), reward, done, info

    def _augment_state(self, state):
        """
        动态归一化角速度，并扩充状态以包含侧风和湍流信息。
        归一化处理：假定侧风最大速度为10 m/s，湍流最大值为2。
        """
        # 对角速度部分（索引24~26）进行动态归一化
        """
        omega = state[24:27]
        self.omega_window = np.roll(self.omega_window, shift=-1, axis=0)
        self.omega_window[-1] = np.abs(omega)
        self.current_omega_max = np.max(self.omega_window, axis=0) + 1e-5  # 防止除零
        normalized_omega = np.clip(omega / self.current_omega_max, -1, 1)
        """

        # 归一化侧风信息：包括侧风速度比例及其方向（以正弦、余弦表示）
        normalized_crosswind = np.array([
            self.crosswind_speed / 10.0,
            np.sin(np.radians(self.crosswind_direction)),
            np.cos(np.radians(self.crosswind_direction))
        ])
        # 归一化湍流强度：简单线性归一化（0~2）
        normalized_turbulence = np.array([self.turbulence_intensity / 2.0] * 3)

        # 将扩展信息追加至原状态中
        augmented_state = np.concatenate([state, normalized_crosswind, normalized_turbulence])
        return augmented_state

    def _apply_environment_disturbance(self, state):
        """
        应用侧风和湍流扰动到状态中以模拟真实飞行中的外界影响。
        侧风扰动：主要影响航向（state[2]，单位：度），根据侧风速度和方向计算偏移。
        湍流扰动：在角速度（state[24~26]）及空速（state[1]）上添加随机噪声。
        """
        # 计算侧风对航向的偏移量：
        crosswind_drift = self.crosswind_speed * 0.1 * np.sin(np.radians(self.crosswind_direction - state[2]))
        state[2] += crosswind_drift

        # 湍流扰动：在角速度上添加随机噪声
        turbulence_noise = np.random.normal(0, self.turbulence_intensity, size=3)
        state[24:27] += turbulence_noise

        # 湍流引起空速的轻微波动
        ias_noise = np.random.normal(0, self.turbulence_intensity * 0.5)
        state[1] += ias_noise

        return state

    def _calculate_balanced_reward(self, state):
        """
        计算平衡多目标的奖励函数：考虑稳定性（角速度、姿态），空速、航点距离及高度，
        同时引入湍流干扰下的额外惩罚，以激励代理在强扰动下保持稳定飞行。
        """
        reward = 0.0
        omega = state[24:27]
        yaw = state[2]
        pitch = state[3]
        roll = state[4]
        speed = state[1]
        aoa = state[11]
        altitude = state[32]

        # 1. 角速度稳定性奖励（指数衰减），鼓励飞机角速度尽量低（即保持平稳）
        omega_penalty = np.sum(50 * np.exp(-0.5 * np.abs(omega)))
        reward += self.reward_weights['stability'] * omega_penalty

        # 新增：湍流惩罚——湍流强度与姿态波动的乘积，体现湍流下机体姿态波动较大时对飞行稳定性的负面影响
        turbulence_penalty = self.turbulence_intensity * (np.abs(pitch) + np.abs(roll))
        reward -= self.reward_weights['stability'] * 0.5 * turbulence_penalty

        # 2. 滚转角惩罚（二次惩罚，以鼓励机体维持较小的姿态偏离，稳定飞行）
        roll_penalty = 0.3 * (roll ** 2)
        reward -= self.reward_weights['stability'] * roll_penalty

        # 3. 速度和高度奖励，俯仰角惩罚
        if self.is_on_ground:
            speed_reward = 80 * np.tanh(speed / TARGET_VA)
            reward += self.reward_weights['speed'] * speed_reward  # 鼓励尽快加速
            if altitude > 50 and speed > self.stall_speed:
                altitude_reward = 400 + 20 * (altitude - 50)
                reward += self.reward_weights['altitude'] * altitude_reward

            if pitch > 20:
                pitch_penalty = 0.2 * ((pitch - 20) ** 2)
            else:
                pitch_penalty = 0  # 不大于20不惩罚
            reward -= self.reward_weights['stability'] * pitch_penalty

        else:
            # 空中阶段：利用当前与目标航点的距离奖励
            dist = self.last_distance
            distance_reward = 1000 * np.log(1 + 1 / (dist + 1e-5))
            reward += self.reward_weights['distance'] * distance_reward  # 鼓励尽快接近航点
            # alt_error = abs(altitude - self.last_altitude)
            # reward += 50 * (1 - alt_error / 100.0)  # 鼓励保持稳定爬升，避免过大高度波动

            # 攻角奖励
            if 10 < aoa < 15:
                reward += 50 * self.reward_weights['aoa']
            else:
                reward += -100 * abs(aoa - 7) * self.reward_weights['aoa']

            # 俯仰角
            if 10 < pitch < 15:
                pitch_penalty = -5 * (pitch ** 2)
            else:
                # 对实际与目标的偏差采用二次惩罚（偏差越大，惩罚越重）
                pitch_penalty = 10 * (pitch ** 2)
            reward -= self.reward_weights['stability'] * pitch_penalty

            # 新增：空中阶段爬升率与垂直速度奖励机制
            # 计算爬升率（假定每步时间间隔为1秒）
            climb_rate = altitude - self.last_altitude  # 单位 m/s
            ideal_climb_rate = 180.0  # 设定理想爬升率为 180 m/s
            if climb_rate >= 0:
                vertical_reward = 50 * (1 - abs(climb_rate - ideal_climb_rate) / ideal_climb_rate)
            else:
                vertical_reward = -1000
            reward += vertical_reward

            target_yaw = 0.0  # 设定期望的偏航为0°
            yaw_error = abs(yaw - target_yaw)
            yaw_penalty = 0.2 * (yaw_error ** 2)
            reward -= self.reward_weights['stability'] * yaw_penalty

        # 4. 时间惩罚（分阶段，平方根形式减缓累积影响）
        time_penalty = 0.5 if self.is_on_ground else 1.5
        reward -= time_penalty * np.sqrt(self.steps)

        return reward

    def _check_dangerous_state(self, state):
        """
        检查危险状态：包括失速、过大姿态偏差、触地反弹等情况，
        以及因侧风或湍流导致的不安全状态。
        """
        if state[1] < self.stall_speed and state[32] > 50:
            return True
        if abs(state[3]) > 45 or abs(state[4]) > 60:
            return True
        if state[32] < (self.last_altitude - 5) and self.steps > 100:
            return True

        # 新增：对偏航的判断，只有当偏航超过60°时认为危险
        if abs(state[2]) > 60:
            return True

        return False

    def _take_action(self, action):
        """
        执行控制动作：在地面阶段采用严格的俯仰限制，
        空中阶段则对控制动作进行主动阻尼，
        并调用基础环境中定义的动作执行函数。
        """
        """
        if self.is_on_ground:
            action[1] = np.clip(action[1], 0, 0.3)
        super()._take_action(action)
        if not self.is_on_ground:
            data = get_dcs_data()
            if data:
                damp_gain = 0.8 + 0.2 * np.tanh(self.steps / 1e4)
                damped_omega = [
                    data["Omega"][0] * damp_gain,
                    data["Omega"][1] * damp_gain,
                    data["Omega"][2] * 0.95
                ]
                self._extract_state({"Omega": damped_omega})
        """
        super()._take_action(action)
